#ifndef _huff_h_
#define _huff_h_

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <sys/stat.h>

#define MAX_MEM 12288 //stopping at 12k
#define BLOCK_WIDTH 16
#define NaB 257

typedef struct Node{
  int value;
  unsigned int prob;
  unsigned int length;
  struct Node *left;
  struct Node *right;
  struct Node *next;
}Node;

typedef struct Element{
  char value;
  unsigned int pos;
  struct Element* next;
}Element;

int huffencode(unsigned char* dst, unsigned char* src, int length, int *comp_size, unsigned int bwi, unsigned int br);
int huffdecode(unsigned char* dst, unsigned char* src, unsigned int *bwi, unsigned int *br);
void setcodes(Node *root, unsigned int l, unsigned long code, unsigned char* dst, int pos, int* dstloc);
Node* findcode(Node *root, char c);

void swap(Node* a, Node* b);
void sort(Node* arr, int beg, int end);

__global__ void initializeAscii(Node *NODE_DEVICE){
  int tx = threadIdx.x;
  int ty = threadIdx.y;
  int index = tx*BLOCK_WIDTH + ty;

  NODE_DEVICE[index].value = index;
  NODE_DEVICE[index].prob = 0;
  
  __syncthreads();
}

/*__global__ void computeProbability(Node *NODE_DEVICE, unsigned char *STRING_DEVICE, int length){
  int localIndex = threadIdx.x*blockDim.x + threadIdx.y;
  int globalIndex = blockIdx.x*blockDim.x*blockDim.y + localIndex;
  int offset;
  int i, j;

  __shared__ int counts[2*BLOCK_WIDTH*BLOCK_WIDTH];
  //__shared__ unsigned char string[2*BLOCK_WIDTH*BLOCK_WIDTH];
  //string[2*localIndex] = STRING_DEVICE[2*globalIndex];
  //string[2*localIndex+1] = STRING_DEVICE[2*globalIndex+1];

  for(i = 0; i < 256; i++){
    __syncthreads();
    memset(counts+2*localIndex*sizeof(int), 0, 2*sizeof(int));
    offset = 1;

    if(STRING_DEVICE[2*globalIndex] == (unsigned char)i && (2*globalIndex) < length)
      counts[2*localIndex] = 1;
    else
      counts[2*localIndex] = 0;
    if(STRING_DEVICE[2*globalIndex+1] == (unsigned char)i && (2*globalIndex+1) < length)
      counts[2*localIndex+1] = 1;
    else
      counts[2*localIndex+1] = 0;

    for(j = (2*blockDim.x*blockDim.y)>> 1; j > 0; j = j >> 1){
      __syncthreads();
      if(localIndex < j){
        int ai = offset*(2*localIndex+1)-1;
        int bi = offset*(2*localIndex+2)-1;
        if(ai < length)
          counts[bi] += counts[ai];
        else
          counts[bi] = counts[ai];
      }
      offset *= 2;
    }

    offset = 1;
    __syncthreads();    
    if(localIndex == 0)
      accum[blockIdx.x] = counts[(2*blockDim.x*blockDim.y) - 1];
    __syncthreads();
    if(blockIdx.x == 0 && localIndex < 64){
      if(localIndex < gridDim.x) counts[localIndex] = accum[localIndex];
      else counts[localIndex] = 0;
      __syncthreads();
      for(j = 64>> 1; j > 0; j = j >> 1){
        __syncthreads();
        if(localIndex < j){
          int ai = offset*(2*localIndex+1)-1;
          int bi = offset*(2*localIndex+2)-1;
          if(ai < gridDim.x)
            counts[bi] += counts[ai];
          else
            counts[bi] = counts[ai];
        }
        offset *= 2;
      }
      __syncthreads();
      if(blockIdx.x == 0 && localIndex == 0){
        for(j = 1; j < gridDim.x; j++){
          accum[0] += accum[j];
        }
        NODE_DEVICE[i].prob = accum[0];
      }
    }
  }
}*/

__global__ void computeProbability(unsigned char *STRING_DEVICE, int *ACCUM_DEVICE, int length, int i){
  int localIndex = threadIdx.x*blockDim.x + threadIdx.y;
  int globalIndex = blockIdx.x*blockDim.x*blockDim.y + localIndex;
  int offset;
  int j;

  __shared__ int counts[2*BLOCK_WIDTH*BLOCK_WIDTH];

  __syncthreads();
  memset(counts+2*localIndex*sizeof(int), 0, 2*sizeof(int));
  offset = 1;

  if(STRING_DEVICE[2*globalIndex] == (unsigned char)i && (2*globalIndex) < length)
    counts[2*localIndex] = 1;
  else
    counts[2*localIndex] = 0;
  if(STRING_DEVICE[2*globalIndex+1] == (unsigned char)i && (2*globalIndex+1) < length)
    counts[2*localIndex+1] = 1;
  else
    counts[2*localIndex+1] = 0;

  for(j = (2*blockDim.x*blockDim.y)>> 1; j > 0; j = j >> 1){
    __syncthreads();
    if(localIndex < j){
      int ai = offset*(2*localIndex+1)-1;
      int bi = offset*(2*localIndex+2)-1;
      if(ai < length)
        counts[bi] += counts[ai];
      else
        counts[bi] = counts[ai];
    }
    offset *= 2;
  }

  __syncthreads();    
  if(localIndex == 0)
    ACCUM_DEVICE[blockIdx.x] = counts[(2*blockDim.x*blockDim.y) - 1];
}

__global__ void sortTable(Node *NODE_DEVICE, Node *NODE_OUT_DEVICE){
  int index = threadIdx.x*blockDim.x + threadIdx.y;
  int i, rank = 0;

  __shared__ Node table[256];
  table[index] = NODE_DEVICE[index]; 

  for (i = 0; i < blockDim.x*blockDim.y; ++i) {
    if (index == i)
      continue;
    if(table[i].prob < table[index].prob) rank++;
    else if(table[i].prob == table[index].prob && i < index) rank++;
  }
  NODE_OUT_DEVICE[rank] = table[index];
  __syncthreads();
}

int huffencode(unsigned char* dst, unsigned char* src, int length, int *comp_size, unsigned int bwi, unsigned int br){
  int i, j, index;
  int size;
  Node table[256], *root, *tree = (Node*)malloc(sizeof(Node)); //*go, *first = tree;
  int numNodes = 1;
  int accum[48];
  Node *NODE_DEVICE, *NODE_OUT_DEVICE;
  int *ACCUM_DEVICE;
  unsigned char *STRING_DEVICE;
  tree->value = 0;
  tree->prob = 0;
  tree->left = NULL;
  tree->right = NULL;
  tree->next = NULL;
  
  /**************************************** 
  * CUDA ascii table initialization
  *****************************************/
  size = 256*sizeof(Node);
  cudaMalloc((void**)&NODE_DEVICE, size);
  if(cudaMemcpy(NODE_DEVICE, table, size, cudaMemcpyHostToDevice)!=cudaSuccess)
    fprintf(stderr, "Copy from host to device failed\n");

  dim3 dimBlock(BLOCK_WIDTH, BLOCK_WIDTH, 1);
  dim3 dimGrid(1, 1, 1);

  initializeAscii<<<dimGrid, dimBlock>>>(NODE_DEVICE);

  if(cudaMemcpy(table, NODE_DEVICE, size, cudaMemcpyDeviceToHost)!=cudaSuccess)
    fprintf(stderr, "Copy from device to host failed\n");

  /**************************************** 
  * sequential string character tally
  *****************************************/
  if(length < 256*log(length)){
    if(cudaMemcpy(table, NODE_DEVICE, size, cudaMemcpyDeviceToHost)!=cudaSuccess)
      fprintf(stderr, "Copy from device to host failed\n");
    for(i = 0; i < length; i++)
      (table[src[i]].prob)++;
    if(cudaMemcpy(NODE_DEVICE, table, size, cudaMemcpyHostToDevice)!=cudaSuccess)
      fprintf(stderr, "Copy from host to device failed\n");
  }
  /**************************************** 
  * CUDA string character tally
  *****************************************/
  else{
    int maxlength = 24576;
    cudaMalloc((void**)&STRING_DEVICE, maxlength);
    cudaMalloc((void**)&ACCUM_DEVICE, 48*sizeof(int));
    for(i = 0; i < (length/maxlength + 1); i++){
      int chunklength;
      if(i == (length/maxlength)) chunklength = length%maxlength;
      else chunklength = maxlength;
      if(chunklength == 0) break;
      else if(chunklength < 256*log(chunklength)){
        for(j = i*maxlength; j < length; j++)
          (table[src[j]].prob)++;
      }
      else{
        if(cudaMemcpy(STRING_DEVICE, src+i*maxlength, chunklength, cudaMemcpyHostToDevice)!=cudaSuccess)
          fprintf(stderr, "Copy from host to device failed\n");

        dimGrid.x = (chunklength/(2*BLOCK_WIDTH*BLOCK_WIDTH));
        if(chunklength%(2*BLOCK_WIDTH*BLOCK_WIDTH) != 0) dimGrid.x++;

        for(index = 0; index < 256; index++){
          computeProbability<<<dimGrid, dimBlock>>>(STRING_DEVICE, ACCUM_DEVICE, chunklength, index);

          if(cudaMemcpy(accum, ACCUM_DEVICE, dimGrid.x*sizeof(int), cudaMemcpyDeviceToHost)!=cudaSuccess)
            fprintf(stderr, "Copy from device to host failed\n");

          for(j = 0; j < (int)dimGrid.x; j++)
            table[index].prob += accum[j];
        }
      }
    }
    cudaFree(STRING_DEVICE);
    cudaFree(ACCUM_DEVICE);

    if(cudaMemcpy(NODE_DEVICE, table, size, cudaMemcpyHostToDevice)!=cudaSuccess)
      fprintf(stderr, "Copy from host to device failed\n");
  }

  /**************************************** 
  * CUDA sort
  *****************************************/
  cudaMalloc((void**)&NODE_OUT_DEVICE, size);

  dimGrid.x = 1;

  sortTable<<<dimGrid, dimBlock>>>(NODE_DEVICE, NODE_OUT_DEVICE);

  if(cudaMemcpy(table, NODE_OUT_DEVICE, size, cudaMemcpyDeviceToHost)!=cudaSuccess)
    fprintf(stderr, "Copy from device to host failed\n");

  /**************************************** 
  * sequential tree creation - Non-parallel
  *****************************************/
  i = 0;
  while(table[i].prob == 0) i++;
  for(; i < 256-1; i++){
    Node *l, *r;
    if(table[i].value != NaB){
      tree->value = table[i].value;
      tree->prob = table[i].prob;
      tree->left = NULL;
      tree->right = NULL;
      l = tree;
      tree->next = (Node*)malloc(sizeof(Node));
	  numNodes++;
      tree = (Node*)tree->next;
    }
    else
      l = (Node*)table[i].next;
    if(table[i+1].value != NaB){
      tree->value = table[i+1].value;
      tree->prob = table[i+1].prob;
      tree->left = NULL;
      tree->right = NULL;
      r = tree;
      tree->next = (Node*)malloc(sizeof(Node));
	  numNodes++;
      tree = (Node*)tree->next;
    }
    else
      r = (Node*)table[i+1].next;
    tree->value = NaB;
    tree->prob = l->prob+r->prob;
    tree->right = (Node*)r;
    tree->left = (Node*)l;
    tree->next = NULL;
    if(i < 254){
      tree->next = (Node*)malloc(sizeof(Node));
	  numNodes++;
      table[i+1].value = NaB;
      table[i+1].prob = (l->prob)+(r->prob);
      table[i+1].next = (Node*)tree;
      tree = (Node*)tree->next;
      tree->value = 0;
      tree->prob = 0;
      tree->left = NULL;
      tree->right = NULL;
      tree->next = NULL;  

      /**************************************** 
      * sequential sort
      *****************************************/
      //sort(table, i+1, 256);

      /**************************************** 
      * CUDA sort
      *****************************************/
      if(cudaMemcpy(NODE_DEVICE, table, size, cudaMemcpyHostToDevice)!=cudaSuccess)
        fprintf(stderr, "Copy from host to device failed\n");

        sortTable<<<dimGrid, dimBlock>>>(NODE_DEVICE, NODE_OUT_DEVICE);

      if(cudaMemcpy(table, NODE_OUT_DEVICE, size, cudaMemcpyDeviceToHost)!=cudaSuccess)
        fprintf(stderr, "Copy from device to host failed\n");

    }
    else
      root = tree;
  }
  int dstloc = 0;
  dst[dstloc++] = (unsigned char)(bwi>>24);
  dst[dstloc++] = (unsigned char)(bwi>>16);
  dst[dstloc++] = (unsigned char)(bwi>>8);
  dst[dstloc++] = (unsigned char)(bwi);
  dst[dstloc++] = (unsigned char)(br>>24);
  dst[dstloc++] = (unsigned char)(br>>16);
  dst[dstloc++] = (unsigned char)(br>>8);
  dst[dstloc++] = (unsigned char)(br);

  setcodes(root, 0, 0, dst, 1, &dstloc);

  /*for(go = first; ; go = (Node*)go->next){ 
     printf("%d %d %d %d %d\n", go->value, go->prob, go, go->left, go->right); 
     if(go->next == NULL) break; 
  }*/

  int treelen = dstloc;
  dst[treelen++] = 255;
  dst[treelen++] = 255;
  dst[treelen++] = 255;
  dst[treelen++] = 255;
  dst[treelen++] = 255;

  int bitlen = 0;
  char c = 0;
  for(i = 0; i < length; i++){
    Node* n = findcode(root, src[i]);
    if(n == NULL) printf(".\\cry %d\n", src[i]);
    for(j = 0; j < n->length; j++){
      if((bitlen%8) == 0)
        c = 0;
      char temp = (n->prob >> (n->length  - j - 1)) & 1;
      c = (c << 1) + temp;
      if((bitlen%8) == 7){
        dst[treelen+bitlen/8] = c;
      }
      bitlen++;
    }
  }

  if(c != 0){
    c = c << (8 - (bitlen%8));
    dst[treelen+bitlen/8] = c;
  }
  int temp = bitlen/8 + 1;
  dst[treelen-4] = (unsigned char)(temp>>24);
  dst[treelen-3] = (unsigned char)(temp>>16);
  dst[treelen-2] = (unsigned char)(temp>>8);
  dst[treelen-1] = (unsigned char)(temp);
  *comp_size = treelen + bitlen/8 +1;

  cudaFree(NODE_DEVICE);
  cudaFree(NODE_OUT_DEVICE);
  return 0;
}

int huffdecode(unsigned char* dst, unsigned char* src, unsigned int *bwi, unsigned int *br){
  int i, j, k;
  Element *list = (Element*)malloc(sizeof(Element)), *element;
  element = list;
  *bwi = (src[0]<<24) + (src[1]<<16) + (src[2]<<8) + src[3];
  *br = (src[4]<<24) + (src[5]<<16) + (src[6]<<8) + src[7];
  free(dst);
  dst = (unsigned char*)malloc(*br+1);
  for(i = 8; src[i] != 255; i=i+4){
    element->value = src[i];

    element->pos = 0;
    element->pos += src[i+1]<<16;
    element->pos += src[i+2]<<8;
    element->pos += src[i+3];

    //printf("%d %d\n", element->value, element->pos);
    element->next = (Element*)malloc(sizeof(Element));
    element = (Element*)element->next;
    element->next = NULL;
  }
  int size = (src[i+1]<<24) + (src[i+2]<<16) + (src[i+3]<<8) + src[i+4] + i + 5;
  i = i + 5;
  //printf("%d\n", size);
  unsigned int c = 1;
  k = 0;
  for(; i <= size; i++){
  //printf("%d %d\n", i, size);
    for(j = 7; j >= 0; j--){
      c = (c<<1) + ((src[i]>>j)&1);
      element = list;
      while(element->next != NULL){
        if(element->pos == c){
          dst[k++] = element->value;
          //printf("%d\n", dst[k-1]);
          c = 1;
          break;
        }
        element = (Element*)element->next;
      }
    }
  }
  dst[k] = 0;
  return k;
}

void setcodes(Node *root, unsigned int l, unsigned long code, unsigned char* dst, int pos, int* dstloc){
  if(root->value != NaB){
    root->prob = code;
    root->length = l;
    dst[(*dstloc)++] = root->value;
    dst[(*dstloc)++] = (unsigned char)(pos>>16);
    dst[(*dstloc)++] = (unsigned char)(pos>>8);
    dst[(*dstloc)++] = (unsigned char)pos;
    //printf("%c\t%d\t%d\t%d\t%d\n", root->value, root->value, root->prob, root->length, pos);
  }
  if(root->left != NULL)
    setcodes((Node*)root->left, l+1, code<<1, dst, 2*pos, dstloc);
  if(root->right != NULL)
    setcodes((Node*)root->right, l+1, (code<<1) + 1, dst, 2*pos+1, dstloc);
}

Node* findcode(Node *root, char c){
  Node* n = NULL;
  if(root->value == c){
    //printf("%c\n", c);
    return root;
  }
  if(root->left != NULL){
    n = findcode((Node*)root->left, c);
    if(n!=NULL) return n;
  }
  if(root->right != NULL)
    n = findcode((Node*)root->right, c);
  return n;
}

#endif
